import gradio as gr

from langchain_openai import OpenAIEmbeddings
from langchain.chains import RetrievalQA
from langchain_openai import ChatOpenAI
from langchain_community.vectorstores import FAISS
from difflib import SequenceMatcher
from langchain.output_parsers import CommaSeparatedListOutputParser


def initialize_estates_sales_bot(vector_store_dir: str="real_estates_sale"):
    db = FAISS.load_local(vector_store_dir, OpenAIEmbeddings(),allow_dangerous_deserialization=True)
    llm = ChatOpenAI(model_name="gpt-4", temperature=0.5)
    
    llm.verbose = True
    global ESTATES_SALES_BOT    
    ESTATES_SALES_BOT = RetrievalQA.from_chain_type(llm,
                                           retriever=db.as_retriever(search_type="similarity_score_threshold",
                                                                     search_kwargs={"score_threshold": 0.8}))
    # 返回向量数据库的检索结果
    ESTATES_SALES_BOT.return_source_documents = True
    ESTATES_SALES_BOT.verbose = True

    global estates
    estates = OpenAIEmbeddings().embed_query("房产销售")

    return ESTATES_SALES_BOT

def initialize_car_sales_bot(vector_store_dir: str="real_car_sale"):
    db = FAISS.load_local(vector_store_dir, OpenAIEmbeddings(),allow_dangerous_deserialization=True)
    llm = ChatOpenAI(model_name="gpt-4", temperature=0.5)
    
    llm.verbose = True
    global CAR_SALES_BOT    
    CAR_SALES_BOT = RetrievalQA.from_chain_type(llm,
                                           retriever=db.as_retriever(search_type="similarity_score_threshold",
                                                                     search_kwargs={"score_threshold": 0.8}))
    # 返回向量数据库的检索结果
    CAR_SALES_BOT.return_source_documents = True
    CAR_SALES_BOT.verbose = True

    global cars
    cars = OpenAIEmbeddings().embed_query("汽车销售")

    return CAR_SALES_BOT

# def estates_sales_chat(message, history):
#     print(f"[message]{message}")
#     print(f"[history]{history}")
#     # TODO: 从命令行参数中获取
#     enable_chat = True

#     ans = ESTATES_SALES_BOT.invoke({"query": message})
#     # 如果检索出结果，或者开了大模型聊天模式
#     # 返回 RetrievalQA combine_documents_chain 整合的结果
#     # print(len(ans["source_documents"]), enable_chat)
#     if len(ans["source_documents"]) or enable_chat:
#         print(f"[result]{ans['result']}")
#         print(f"[source_documents]{ans['source_documents']}")
#         return ans["result"]
#     # 否则输出套路话术
#     else:
#         return "这个问题我要问问领导"
    
# def car_sales_chat(message, history):
#     print(f"[message]{message}")
#     print(f"[history]{history}")
#     # TODO: 从命令行参数中获取
#     enable_chat = True

#     ans = CAR_SALES_BOT.invoke({"query": message})
#     # 如果检索出结果，或者开了大模型聊天模式
#     # 返回 RetrievalQA combine_documents_chain 整合的结果
#     # print(len(ans["source_documents"]), enable_chat)
#     if len(ans["source_documents"]) or enable_chat:
#         print(f"[result]{ans['result']}")
#         print(f"[source_documents]{ans['source_documents']}")
#         return ans["result"]
#     # 否则输出套路话术
#     else:
#         return "这个问题我要问问领导"

def sales_chat(message, history):
    print(f"[message]{message}")
    print(f"[history]{history}")
    # TODO: 从命令行参数中获取
    enable_chat = True
    # 如果是房产销售问题
    estates_similarity = SequenceMatcher(None, message, '房产销售').ratio()
    print(f"[estates_similarity]{estates_similarity}")
    car_similarity = SequenceMatcher(None, message, '汽车销售').ratio()
    print(f"[car_similarity]{car_similarity}")
    if estates_similarity >= car_similarity:
        ans = ESTATES_SALES_BOT.invoke({"query": message})
    else:
        ans = CAR_SALES_BOT.invoke({"query": message})
    # 如果检索出结果，或者开了大模型聊天模式
    # 返回 RetrievalQA combine_documents_chain 整合的结果
    # print(len(ans["source_documents"]), enable_chat)
    if len(ans["source_documents"]) or enable_chat:
        print(f"[result]{ans['result']}")
        print(f"[source_documents]{ans['source_documents']}")
        return ans["result"]
    # 否则输出套路话术
    else:
        return "这个问题我要问问领导"

def launch_gradio():
    demo = gr.ChatInterface(
        fn=sales_chat,
        title="销售机器人",
    )
    
    demo.launch(share=False, server_name="0.0.0.0")

if __name__ == "__main__":
    # 初始化房产销售机器人
    initialize_estates_sales_bot()
    # 初始化汽车销售机器人
    initialize_car_sales_bot()
    # 启动 Gradio 服务
    launch_gradio()
